#include "Tokenizer.hpp"

namespace Welp
{
  DfaEdge::DfaEdge() :
    mToState(nullptr)
  {
  }

  void DfaEdge::AddRange(unsigned char start, unsigned char end)
  {
    WelpErrorIf(start > end, "Start must be before or equal to end");
    for (unsigned char i = start; i < end; ++i)
    {
      mValidChars.insert(i);
    }
  }

  DfaState::DfaState() :
    mAcceptingToken(TokenType::Invalid),
    mErrorOnFail(nullptr),
    mName(nullptr)
  {
  }
 
  DfaEdge* DfaState::FindEdge(char c)
  {
    for (size_t i = 0; i < mOutgoing.size(); ++i)
    {
      DfaEdge* edge = mOutgoing[i];

      if (edge->mValidChars.find(c) != edge->mValidChars.end())
      {
        return edge;
      }
    }

    return nullptr;
  }

  DfaState* Dfa::CreateState()
  {
    auto state = new DfaState();
    mAllStates.push_back(state);
    return state;
  }

  Dfa::Dfa()
  {
    mRoot = CreateState();
  }

  Dfa::~Dfa()
  {
    for (size_t i = 0; i < mAllStates.size(); ++i)
    {
      delete mAllStates[i];
    }

    for (size_t i = 0; i < mAllEdges.size(); ++i)
    {
      delete mAllEdges[i];
    }
  }

  DfaEdge* Dfa::CreateOutgoingEdge(DfaState* from, DfaState* to)
  {
    auto edge = new DfaEdge();
    mAllEdges.push_back(edge);
    from->mOutgoing.push_back(edge);
    edge->mToState = to;
    return edge;
  }

  DfaTokenizer::DfaTokenizer()
  {
    // Eof
    {
      auto s1 = mDfa.CreateState();
      s1->mName = "End Of File";
      s1->mAcceptingToken = TokenType::Eof;
      auto rootToS1 = mDfa.CreateOutgoingEdge(mDfa.mRoot, s1);
      rootToS1->mValidChars.insert('\0');
    }

    // Whitespace
    {
      auto s1 = mDfa.CreateState();
      s1->mName = "Whitespace";
      s1->mAcceptingToken = TokenType::Whitespace;
      auto rootToS1 = mDfa.CreateOutgoingEdge(mDfa.mRoot, s1);
      rootToS1->mValidChars.insert('\f');
      rootToS1->mValidChars.insert('\n');
      rootToS1->mValidChars.insert('\r');
      rootToS1->mValidChars.insert('\t');
      rootToS1->mValidChars.insert('\v');
      rootToS1->mValidChars.insert(' ');

      auto s1ToS1 = mDfa.CreateOutgoingEdge(s1, s1);
      s1ToS1->mValidChars.insert('\f');
      s1ToS1->mValidChars.insert('\n');
      s1ToS1->mValidChars.insert('\r');
      s1ToS1->mValidChars.insert('\t');
      s1ToS1->mValidChars.insert('\v');
      s1ToS1->mValidChars.insert(' ');
    }

    // String
    {
      auto s1 = mDfa.CreateState();
      s1->mName = "String";
      auto rootToS1 = mDfa.CreateOutgoingEdge(mDfa.mRoot, s1);
      rootToS1->mValidChars.insert('"');

      auto s2 = mDfa.CreateState();
      s2->mErrorOnFail = "Invalid escape sequence";
      auto s1ToS2 = mDfa.CreateOutgoingEdge(s1, s2);
      s1ToS2->mValidChars.insert('\\');

      auto s2ToS1 = mDfa.CreateOutgoingEdge(s2, s1);
      s2ToS1->mValidChars.insert('n');
      s2ToS1->mValidChars.insert('r');
      s2ToS1->mValidChars.insert('t');
      s2ToS1->mValidChars.insert('\\');
      s2ToS1->mValidChars.insert('"');

      auto s3 = mDfa.CreateState();
      s3->mAcceptingToken = TokenType::String;
      auto s1ToS3 = mDfa.CreateOutgoingEdge(s1, s3);
      s1ToS3->mValidChars.insert('"');

      auto s1ToS1 = mDfa.CreateOutgoingEdge(s1, s1);
      // Allow every character to appear in the string, except another " or \ (those can appear via escapes)
      s1ToS1->AddRange((char)1, (char)255);
      s1ToS1->mValidChars.erase('"');
      s1ToS1->mValidChars.erase('\\');
    }

    // Identifier
    {
      auto s1 = mDfa.CreateState();
      s1->mName = "Identifier";
      s1->mAcceptingToken = TokenType::Identifier;
      auto rootToS1 = mDfa.CreateOutgoingEdge(mDfa.mRoot, s1);
      rootToS1->AddRange('a', 'z');
      rootToS1->AddRange('A', 'Z');
      rootToS1->mValidChars.insert('_');

      auto s1ToS1 = mDfa.CreateOutgoingEdge(s1, s1);
      s1ToS1->AddRange('a', 'z');
      s1ToS1->AddRange('A', 'Z');
      s1ToS1->AddRange('0', '9');
      s1ToS1->mValidChars.insert('_');
    }
  }

  
  //void DfaTokenizer::Visualize(DfaEdge* edge)
  //{
  //  GraphVisualizer graph;
  //  graph.BeginDiGraph();
  //  graph.AddNode("Root", "", "box");


  //  string lastStateName = "Root";


  //  unordered_set<DfaState*> visitedStates;

  //  auto state = edge->mToState;
  //  string newStateName = Format("%p", state);

  //  for (size_t i = 0; i < state->mOutgoing.size(); ++i)
  //  {
  //    edge = state->mOutgoing[i];

  //    
  //    graph.AddEdge(lastStateName.c_str(), newStateName.c_str());
  //  }

  //  //for (size_t i = 0; i < rootEdge->mValidChars
  //  //rootEdge->mToState

  //  
  //  graph.EndGraph();
  //  graph.Visualize();
  //}

  void DfaTokenizer::RegisterKeyword(const char* keyword)
  {
  }

  void DfaTokenizer::RegisterSymbol(const char* symbol)
  {
    DfaState* state = mDfa.mRoot;

    size_t length = strlen(symbol);
    for (size_t i = 0; i < length; ++i)
    {
      char c = symbol[i];

      DfaEdge* edge = state->FindEdge(c);

      if (edge == nullptr)
      {
      }
      else
      {
      }
    }
  }

  void DfaTokenizer::DefaultErrorCallback(const string& message, const Location& location, void* userData)
  {
    printf("%s(%d): Welp Token Error: %s\n", location.mOrigin.c_str(), location.mLine, message.c_str());
  }

  string EscapeChar(char c)
  {
    switch (c)
    {
      case '\0':
        return "\\0";
      case '\a':
        return "\\a";
      case '\b':
        return "\\b";
      case '\f':
        return "\\f";
      case '\n':
        return "\\n";
      case '\r':
        return "\\r";
      case '\t':
        return "\\t";
      case '\v':
        return "\\v";
    }
    auto str = string("x");
    str[0] = c;
    return str;
  }

  void DfaTokenizer::ReadToken(Token* tokenOut)
  {
    tokenOut->mLocation = mLocation;
    tokenOut->mType = TokenType::Invalid;
    tokenOut->mText.clear();

    size_t lineNumber = mLocation.mLine;
    size_t characterNumber = mLocation.mCharacter;

    if (mLocation.mPosition > mText.size())
      return;

    auto start = mLocation.mPosition;
    auto end = (size_t)-1;

    const char* name = nullptr;
    DfaState* state = mDfa.mRoot;

    for (;;)
    {
      if (name == nullptr && state->mName != nullptr)
      {
        name = state->mName;
      }

      char c = mText.c_str()[mLocation.mPosition];
      ++mLocation.mPosition;

      if (c == '\n')
      {
        ++lineNumber;
        characterNumber = 1;
      }
      else
      {
        ++characterNumber;
      }

      DfaEdge* edge = state->FindEdge(c);

      if (edge != nullptr)
      {
        state = edge->mToState;
      }
      else
      {
        if (tokenOut->mType != TokenType::Invalid)
        {
          mLocation.mPosition = end;
          tokenOut->mText = mText.substr(start, end - start);
          return;
        }
        else
        {
          string preText;
          if (state->mErrorOnFail != nullptr)
          {
            preText = Format("%s. ", state->mErrorOnFail);
          }

          string postText;
          if (name != nullptr)
          {
            postText = Format(" when parsing '%s'", name);
          }

          string fullError = Format
          (
            "%sRead invalid character '%s' (Ascii: %d)%s.",
            preText.c_str(),
            EscapeChar(c).c_str(),
            c,
            postText.c_str()
          );

          mLocation.mPosition = start;

          throw Exception(fullError, ExceptionType::Tokenizer, &mLocation);
        }
      }
      
      if (state->mAcceptingToken != TokenType::Invalid)
      {
        end = mLocation.mPosition;
        tokenOut->mType = state->mAcceptingToken;
        mLocation.mLine = lineNumber;
        mLocation.mCharacter = characterNumber;
      }
    }
  }

  Tokenizer::Tokenizer()
  {
    mPosition = 0;

    mStringToFixedToken["true"]      = TokenType::True;
    mStringToFixedToken["false"]     = TokenType::False;
    mStringToFixedToken["if"]        = TokenType::If;
    mStringToFixedToken["while"]     = TokenType::While;
    mStringToFixedToken["function"]  = TokenType::Function;
    
    mStringToFixedToken[";"]   = TokenType::StatementSeparator;
    mStringToFixedToken[","]   = TokenType::ArgumentSeparator;;
    mStringToFixedToken["("]   = TokenType::BeginGroup;
    mStringToFixedToken[")"]   = TokenType::EndGroup;
    mStringToFixedToken["{"]   = TokenType::BeginScope;
    mStringToFixedToken["}"]   = TokenType::EndScope;
    mStringToFixedToken["["]   = TokenType::BeginIndex;
    mStringToFixedToken["]"]   = TokenType::EndIndex;
    mStringToFixedToken["."]   = TokenType::MemberAccess;
    mStringToFixedToken["+"]   = TokenType::Plus;
    mStringToFixedToken["-"]   = TokenType::Minus;
    mStringToFixedToken["!"]   = TokenType::LogicalNot;
    mStringToFixedToken["~"]   = TokenType::BitwiseNot;
    mStringToFixedToken["++"]  = TokenType::Increment;
    mStringToFixedToken["--"]  = TokenType::Decrement;
    mStringToFixedToken["*"]   = TokenType::Multiply;
    mStringToFixedToken["/"]   = TokenType::Divide;
    mStringToFixedToken["%"]   = TokenType::Modulo;
    mStringToFixedToken["^"]   = TokenType::Exponent;
    mStringToFixedToken["<<"]  = TokenType::BitshiftLeft;
    mStringToFixedToken[">>"]  = TokenType::BitshiftRight;
    mStringToFixedToken["<"]   = TokenType::LessThan;
    mStringToFixedToken["<="]  = TokenType::LessThanOrEqualTo;
    mStringToFixedToken[">"]   = TokenType::GreaterThan;
    mStringToFixedToken[">="]  = TokenType::GreaterThanOrEqualTo;
    mStringToFixedToken["=="]  = TokenType::EqualTo;
    mStringToFixedToken["!="]  = TokenType::NotEqualTo;
    mStringToFixedToken["&"]   = TokenType::BitwiseAnd;
    mStringToFixedToken["$"]   = TokenType::BitwiseXor;
    mStringToFixedToken["|"]   = TokenType::BitwiseOr;
    mStringToFixedToken["&&"]  = TokenType::LogicalAnd;
    mStringToFixedToken["||"]  = TokenType::LogicalOr;
    mStringToFixedToken["=="]  = TokenType::Assignment;
    mStringToFixedToken["+="]  = TokenType::PlusAssignment;
    mStringToFixedToken["-="]  = TokenType::MinusAssignment;
    mStringToFixedToken["*="]  = TokenType::MultiplyAssignment;
    mStringToFixedToken["/="]  = TokenType::DivideAssignment;
    mStringToFixedToken["%="]  = TokenType::ModuloAssignment;
    mStringToFixedToken["^="]  = TokenType::ExponentAssignment;
    mStringToFixedToken["<<="] = TokenType::BitshiftLeftAssignment;
    mStringToFixedToken[">>="] = TokenType::BitshiftRightAssignment;
    mStringToFixedToken["&="]  = TokenType::BitwiseAndAssignment;
    mStringToFixedToken["$="]  = TokenType::BitwiseXorAssignment;
    mStringToFixedToken["|="]  = TokenType::BitwiseOrAssignment;

    auto fixedTokens = mStringToFixedToken.begin();
    while (fixedTokens != mStringToFixedToken.end())
    {
      mTokenToString[fixedTokens->second] = fixedTokens->first;
      ++fixedTokens;
    }

    mTokenToString[TokenType::Invalid]    = "<invalid>";
    mTokenToString[TokenType::Number]     = "<number literal>";
    mTokenToString[TokenType::String]     = "<string literal>";
    mTokenToString[TokenType::Identifier] = "<identifier>";
    mTokenToString[TokenType::Symbol]     = "<symbol>";
  }

  void Tokenizer::ReadAndClassifyToken(Token* tokenOut)
  {
    ReadToken(tokenOut);

    auto foundToken = mStringToFixedToken.find(tokenOut->mText);

    if (foundToken != mStringToFixedToken.end())
    {
      tokenOut->mType = foundToken->second;
    }

    // If the token still thinks it's a symbol it may be an issue of concatenated symbols, eg >=( in 5>=(3+4)
    if (tokenOut->mType == TokenType::Symbol)
    {
      size_t longestSymbolMatch = 0;

      // This part obviously takes a lot more time, but it's effective
      // We just look for the longest substring of a symbol that we can match at the beginning
      auto it = mStringToFixedToken.begin();
      while (it != mStringToFixedToken.end())
      {
        const string& testingSymbol = it->first;

        // If our unknown symbol is shorter than the one we're testing, skip it!
        if (testingSymbol.size() > tokenOut->mText.size())
        {
          ++it;
          continue;
        }

        // If we found a match...
        if (testingSymbol.size() > longestSymbolMatch && memcmp(testingSymbol.c_str(), tokenOut->mText.c_str(), testingSymbol.size()) == 0)
        {
          longestSymbolMatch = testingSymbol.size();

          // Backup the tokenizer to the next symbol
          mPosition -= tokenOut->mText.size();
          mPosition += testingSymbol.size();

          tokenOut->mText = testingSymbol;
          tokenOut->mType = it->second;
        }
        ++it;
      }
    }
  }

  void Tokenizer::ReadToken(Token* tokenOut)
  {
    tokenOut->mType = TokenType::Invalid;
    tokenOut->mText.clear();

    const char* text = mText.c_str();

    // Eat leading whitespace...
    while (mPosition < mText.size())
    {
      char c = text[mPosition];

      if (!isspace(c))
      {
        break;
      }

      ++mPosition;
    }

    size_t tokenStart = mPosition;

    char deciderChar = text[mPosition];

    if (deciderChar == '\0')
    {
      // Returning an invalid token signals that we're done!
      return;
    }
    else if (deciderChar == '"')
    {
      tokenOut->mType = TokenType::String;
    }
    else if (isalpha(deciderChar))
    {
      tokenOut->mType = TokenType::Identifier;
    }
    else if (isdigit(deciderChar))
    {
      tokenOut->mType = TokenType::Number;
    }
    else if (deciderChar == '/')
    {
      tokenOut->mType = TokenType::Comment;
    }
    else
    {
      tokenOut->mType = TokenType::Symbol;
    }

    ++mPosition;

    // Read the rest of the token
    while (mPosition <= mText.size())
    {
      char c = text[mPosition];

      if (c == '\0')
      {
        break;
      }

      switch (tokenOut->mType)
      {
        case TokenType::String:
        {
          // Problem, this doesn't allow us to escape quotes...
          if (c == '"')
          {
            // Most other token types read one past itself, but a string knows
            // when it's done by hitting the ending ", which means we need to move the position to the next character
            // so that the next token read doesn't think it's a string!
            ++mPosition;
            tokenOut->mText = mText.substr(tokenStart, mPosition - tokenStart);
            return;
          }
          break;
        }

        case TokenType::Comment:
        {
          // Comments stop when we hit a newline
          if (c == '\n' || c == '\r')
          {
            // We don't need to move the token ahead, whitespace will be eaten on the next token read
            tokenOut->mText = mText.substr(tokenStart, mPosition - tokenStart);
            return;
          }
          break;
        }

        case TokenType::Identifier:
        {
          // Once we hit a non-alpha numeric character, we're done
          if (!isalnum(c))
          {
            tokenOut->mText = mText.substr(tokenStart, mPosition - tokenStart);
            return;
          }
          break;
        }

        case TokenType::Number:
        {
          // We allow digits and a '.' (eventually we'll support exponentials)
          // This has an error case, what if I did 3.5.6?
          if (!(isdigit(c) || c == '.'))
          {
            tokenOut->mText = mText.substr(tokenStart, mPosition - tokenStart);
            return;
          }
          break;
        }

        case TokenType::Symbol:
        {
          // We allow digits and a '.' (eventually we'll support exponentials)
          // This has an error case, what if I did 3.5.6?
          if (isalnum(c) || isspace(c) || c == '"')
          {
            tokenOut->mText = mText.substr(tokenStart, mPosition - tokenStart);
            return;
          }
          break;
        }
      }

      ++mPosition;
    }

    // The only way we get here is if we hit the end of the text
    tokenOut->mText = mText.substr(tokenStart, mPosition - tokenStart);
  }
}
